package com;


import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.OAEPParameterSpec;
import javax.crypto.spec.PSource;
import javax.crypto.spec.SecretKeySpec;

import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.InvalidAlgorithmParameterException;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.MGF1ParameterSpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.security.InvalidAlgorithmParameterException;


public class ClientEncryptionExample {
    /**
     * 生成一个12字节（96位）的随机IV（初始化向量）
     * @return 12字节的IV数组
     */
    public static byte[] generateIvTag(int len) {
        byte[] iv = new byte[len]; // 12字节 = 96位
        SecureRandom secureRandom = new SecureRandom();
        secureRandom.nextBytes(iv);
        return iv;
    }

    public static void main(String[] args) throws Exception {
        // 1. 生成随机AES密钥
        KeyGenerator aesKeyGen = KeyGenerator.getInstance("AES");
        aesKeyGen.init(256);
        SecretKey aesKey = aesKeyGen.generateKey();
        byte[] iv = generateIvTag(12); //new byte[12]; // 实际应生成随机IV并发送到服务端

        // 2. 使用AES加密数据
        String body = aesEncryPtion(aesKey, "{\"title\":\"Java\"}", iv);

        // 3. 使用服务端公钥加密AES密钥
        String b64Key = rsaEncryPtion(aesKey);

        String b64IV = Base64.getEncoder().encodeToString(iv);
        // String body = Base64.getEncoder().encodeToString(encryptedData);
        System.out.println("发送内容：");
        System.out.println("key:" + b64Key);
        System.out.println("iv:" + b64IV);
        System.out.println("body:" + body);
        // 4. 发送请求（示例使用HttpClient）
        HttpClient client = HttpClient.newHttpClient();
        HttpRequest request = HttpRequest.newBuilder()
                .header("Content-Type", "application/json")
                .uri(URI.create("http://localhost:8080/dataShare/testEncryption"))
                .header("X-Encrypted-AES-Key", b64Key)
                .header("X-Encrypted-AES-IV", b64IV)
//                .header("X-Encrypted-Data", Base64.getEncoder().encodeToString(encryptedData))
                .POST(HttpRequest.BodyPublishers.ofString(body)) // 实际发送加密数据
                .build();
        try {
            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
            System.out.println("返回内容：");
            System.out.println(response.body());
            String responseIV = response.headers().firstValue("X-Encrypted-AES-IV").orElse(null);

            aesDecryPtion(aesKey.getEncoded(), response.body(), responseIV);
        } catch (IOException e) {
            e.printStackTrace();
            System.err.println("错误详情: " + e.getMessage());
        }
    }

    // RSA 加密
    private static String rsaEncryPtion(SecretKey aesKey){
        OAEPParameterSpec oaepParams = new OAEPParameterSpec(
            "SHA-256",
            "MGF1",
            MGF1ParameterSpec.SHA256,
            PSource.PSpecified.DEFAULT
        );

        PublicKey serverPublicKey;
        try {
            serverPublicKey = loadPublicKey("../key/public_key.pem");
            Cipher rsaCipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
            rsaCipher.init(Cipher.ENCRYPT_MODE, serverPublicKey, oaepParams);
            byte[] encryptedAesKey = rsaCipher.doFinal(aesKey.getEncoded());
            return Base64.getEncoder().encodeToString(encryptedAesKey);
        } catch (Exception e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
            return "";
        }

    }

    // AES 加密
    private static String aesEncryPtion(SecretKey aesKey, String content, byte[] iv){
        try{
            Cipher aesCipher = Cipher.getInstance("AES/GCM/NoPadding");
            GCMParameterSpec gcmSpec = new GCMParameterSpec(128, iv);
            aesCipher.init(Cipher.ENCRYPT_MODE, aesKey, gcmSpec);
            // String plainText = "{\"title\":\"Java\"}";
            byte[] encryptedData = aesCipher.doFinal(content.getBytes(StandardCharsets.UTF_8));
            return Base64.getEncoder().encodeToString(encryptedData);
        }catch(Exception e){
            return "";
        }

    }

    // AES 解密
    private static String aesDecryPtion(byte[] aesKey, String content, String base64IV){
        
        try {
            Cipher aesCipher1 = Cipher.getInstance("AES/GCM/NoPadding");
            SecretKey aesKey1 = new SecretKeySpec(aesKey, "AES");
            GCMParameterSpec gcmSpec1 = new GCMParameterSpec(128, Base64.getDecoder().decode(base64IV));
            aesCipher1.init(Cipher.DECRYPT_MODE, aesKey1, gcmSpec1);
            String decryptedData = new String(aesCipher1.doFinal(Base64.getDecoder().decode(content)), StandardCharsets.UTF_8);
            System.out.println(decryptedData);
            return decryptedData;
        } catch (Exception e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
            return "";
        } 

    }
    
    // 加载RSA密钥
    private static PublicKey loadPublicKey(String filename) throws Exception {
        String publicKeyContent = new String(Files.readAllBytes(Paths.get(filename)))
                .replaceAll("\\n", "")
                .replaceAll("\\r", "")
                .replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "");
        byte[] decoded = Base64.getDecoder().decode(publicKeyContent);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(decoded);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePublic(keySpec);
    }
}